package com.jsz.yun.ai.field.util.ttl;

import java.util.*;

/**
 * @author liuhognru
 * @date 2023/12/15
 * @description 用于主线程与线程池线程中本地变量的传递
 */
public class TransferableThreadLocal<T> extends InheritableThreadLocal<T> {

    /**
     * 存储当前线程与多个ThreadLocal对象的映射关系
     */
    private static final ThreadLocal<Map<ThreadLocal<?>, Object>> THREAD_LOCAL_MAP = ThreadLocal.withInitial(WeakHashMap::new);

    @Override
    protected T initialValue() {
        T value = super.initialValue();
        Map<ThreadLocal<?>, Object> map = THREAD_LOCAL_MAP.get();
        map.put(this, value);
        return value;
    }

    @Override
    public void set(T value) {
        super.set(value);
        Map<ThreadLocal<?>, Object> map = THREAD_LOCAL_MAP.get();
        map.put(this, value);
    }

    @Override
    public void remove() {
        super.remove();
        Map<ThreadLocal<?>, Object> map = THREAD_LOCAL_MAP.get();
        map.remove(this);
    }

    public static Map<ThreadLocal<?>, Object> getTlMap() {
        return THREAD_LOCAL_MAP.get();
    }

    /**
     * 清理与当前线程绑定的ThreadLocal，仅用于异步线程执行时调用
     * 考虑到ThreadPoolExecutor.CallerRunsPolicy()的拒绝策略是当前线程执行，
     * 如果是创建当前runnable任务的线程则不清理，否则若任务执行完后还有后续代码
     * 要获取ThreadLocal的值，下次拿到的值会为null
     */
    public static void clear(Thread mainThread) {
        if (!Objects.equals(Thread.currentThread(), mainThread)) {
            clear();
        }
    }

    /**
     * 清理与当前线程绑定的ThreadLocal，仅用于主线程执行时调用
     */
    public static void clear() {
        Map<ThreadLocal<?>, Object> map = THREAD_LOCAL_MAP.get();
        //防止迭代器遍历过程中删除元素导致ConcurrentModificationException
        List<Map.Entry<ThreadLocal<?>, Object>> list = new ArrayList<>(map.entrySet());
        list.forEach(entry -> entry.getKey().remove());
        map.clear();
        THREAD_LOCAL_MAP.remove();
    }
}
